package edu.vt.ece.util;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Vector;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantLock;

import edu.vt.ece.main.Tester;

public class AVLTree{

	protected Node root;
	public ReentrantLock treeLock; //this is for the course-grained locking mechanism
	protected Vector<Integer> itemsAdded = new Vector<Integer>();
	public int extraBalance = 0;
	protected AtomicInteger rebalancingRoot = new AtomicInteger(-1);
	protected AtomicInteger rebalancingLeft = new AtomicInteger(-1);

	public AVLTree() {
		treeLock = new ReentrantLock();
	}

	public AVLTree(int key) {
		root = new Node(key);
		itemsAdded.add(key);
		treeLock = new ReentrantLock();
	}

	public Node getRoot() {
		return root;
	}

	public void setRoot(Node node) {
		root = node;
	}
	
	public void getTreeLock()
	{
		treeLock.lock();
	}

	public void releaseTreeLock()
	{
		treeLock.unlock();
	}

	/**
	 * Adds the item into the tree. Duplicates are not allowed.
	 *
	 * @param item
	 * @return true if the item was added successfully, false if it was already
	 *         in the tree
	 */
	public boolean add(int item) {

		Node newNode = new Node(item);
		if (root == null)
		{
			//new tree, root is the first node, return immediately
			root = newNode;
			return true;
		}

		Node test = root; // this is the item that we are going to compare with item

		while (true)
		{
			while (item < test.getKey()) {
				// go left
				if (test.getLeft() != null) {
					test = test.getLeft();
				} else {
					break;
				}
			}
			while (item > test.getKey()) {
				// go right
				if (test.getRight() != null) {
					test = test.getRight();
				} else {
					break;
				}
			}
			if ((item < test.getKey() && test.getLeft() == null) || (item > test.getKey() && test.getRight() == null))
			{
				break;
			}
			if (item == test.getKey())
			{
				return false;
			}
		}

		if (item == test.getKey()) {
			return false; // duplicate key
		}
		itemsAdded.add(item); //only want unique items
		// we have gone as far as we can, now we insert the item


		if (item < test.getKey()) {
			test.setLeft(newNode);
			newNode.setParent(test);
		} else {
			// item > test.getKey()
			test.setRight(newNode);
			newNode.setParent(test);
		}
		rebalanceAndSetRoot(newNode);
		return true;
	}

	private void rebalanceAndSetRoot(Node node) {

		while(node.getParent()!=null) {
			Node parentNode = node.getParent();
			
			if(Math.abs(node.getBalanceFactor()) > 1) rebalance(node);
			node = parentNode;
		}
		if (Math.abs(node.getBalanceFactor()) > 1)
		{
			rebalance(node);
		}

		//sets root to be the parent-most parent
		while (node.getParent() != null) {
			node = node.getParent();
		}
		root = node;
		isValidTree(root);
	}

	/*
	 * Removes the item from the tree.
	 */
	public boolean remove(int key) {
		Node test = contains(key);
		if (test == null)
		{
			return false;
		}
		itemsAdded.add(0 - key); //the negative items indicate those that have been removed
		Node tmpNode = removeNode(test, null);

		if(tmpNode!=null) rebalanceAndSetRoot(tmpNode);

		return true;
	}

	private Node removeNode(Node test, Node replacementNode) {

		//determine test's relationship to parent, store in isLeft
		boolean isRoot = false;
		boolean isLeft = false;
		if(test.getParent() == null) {
			isRoot = true;
		} else {
			if(test.getParent().getLeft() == test) {
				//test is parent's left
				isLeft = true;
			}
		}

		// if test is leaf node
		if(test.getLeft() == null && test.getRight() == null) {
			if(isRoot) {
				root = null;
				return root;
			}

			if(isLeft) test.getParent().setLeft(null);
			else test.getParent().setRight(null);
			Node parent = test.getParent();
			test.setParent(null);

			if(replacementNode!=null) replacementNode.setKey(test.getKey());
			return parent;
		}

		// if test has one leaf
		if(test.getLeft() == null && test.getRight() != null) {
			if(isRoot) {
				root = test.getRight();
			}
			else {
				if(isLeft) test.getParent().setLeft(test.getRight());
				else test.getParent().setRight(test.getRight());
			}
			test.getRight().setParent(test.getParent());
			Node parent = test.getParent();
			test.setParent(null);

			if(replacementNode!=null) replacementNode.setKey(test.getKey());
			return parent;
		} else if(test.getLeft() != null && test.getRight() == null) {
			if(isRoot) {
				root = test.getLeft();
			}
			else {
				if(isLeft) {
					test.getParent().setLeft(test.getLeft());
				} else {
					test.getParent().setRight(test.getLeft());
				}
			}
			test.getLeft().setParent(test.getParent());
			Node parent = test.getParent();
			test.setParent(null);

			if(replacementNode!=null) replacementNode.setKey(test.getKey());
			return parent;
		}


		// if there are two branches
		Node tmpNode = null;
		if(test.getLeft() != null) {
			tmpNode = test.getLeft();
			while(tmpNode.getRight()!=null) tmpNode = tmpNode.getRight();
		}
		else if(test.getRight() != null) {
			tmpNode = test.getRight();
			while(tmpNode.getRight()!=null) tmpNode = tmpNode.getLeft();
		}
		else System.err.println("Should never be here.");

		return removeNode(tmpNode, test);
	}

	/**
	 * Tests whether or not a node with the specified key exists in the tree
	 * If it does, the node with that value is returned, null otherwise
	 * @param key they value to search the tree for
	 * @return the node with the specified key; false otherwise
	 */
	public Node contains(int key) {

		if(root==null) return null;

		Node test = root; // this is the item that we are going to compare with item

		while(key!=test.getKey()) {
			if(test.getLeft()==null && test.getRight()==null) break;
			if(key < test.getKey()) {
				if(test.getLeft()!=null) test = test.getLeft();
				else
					return null; //item is smaller than test but is not in the tree
			}
			else if(key > test.getKey()) {
				if(test.getRight()!=null) test = test.getRight();
				else return null;
			}
		}
		if(key == test.getKey()) return test;
		return null;
	}

	/**
	 * Rebalance the tree at the given node by determining which type of rotation is needed to bring the tree back
	 * into balance
	 * @param node the node from which to balance the tree
	 */
	protected void rebalance(Node node)
    {
	    getTreeLock();
	    try
	    {
		if(node.getKey() == root.getKey()) {
			rebalancingRoot.set(1);
		} else if(node.getKey() < root.getKey()) {
			rebalancingLeft.set(1);
			rebalancingRoot.set(0);
		} else{
			rebalancingLeft.set(0);
			rebalancingRoot.set(0);
		}
		
		int balanceFactor = node.getBalanceFactor();
		if (Math.abs(balanceFactor) > 1)
        {
			// need to rebalance the tree
			if (balanceFactor == 2) {
				// need to do either a right or a leftRight rotation
				switch (node.getLeft().getBalanceFactor()) {
				case -1:
					leftRightRotation(node);
					break;
				case 1:
				case 0:
					singleRightRotation(node);
					break;
				default:
					rebalance(node.getLeft());
				}
			}
			else if (balanceFactor == -2)
			{
				// need to do either a left or a rightLeft rotation
				switch (node.getRight().getBalanceFactor()) {
				case -1:
					singleLeftRotation(node);
					break;
				case 1:
				case 0:
					rightLeftRotation(node);
					break;
				default:
					rebalance(node.getRight());
				}
			}
	        else
			{
	        	if(balanceFactor<0) {
	        		rebalance(node.getRight());	        		
	        	}
	        	else {
	        		rebalance(node.getLeft());
	        	}
			}
		}
	    }
	    finally
	    {
		    releaseTreeLock();
	    }
	}

	/**
	 * The tree structure looks like a linked list from node with two children
	 * to the right and no left children This is the mirrored operation of
	 * singleRightRotation
	 *
	 * Need to move node to left child of node's right child and make node's
	 * parent point to node's right
	 *
	 * @param node - the node about which to perform the rotation
	 */
	private void singleLeftRotation(Node node) {
		Node right = node.getRight(); // this node will be the new "parent"
		if (node.getParent() != null && node.getKey() < node.getParent().getKey())
		{
			// node is the left child of the parent
			node.getParent().setLeft(right);
		}
		else if (node.getParent() != null)
		{
			// node is the parent's right child, set that reference instead
			node.getParent().setRight(right);
		}
		//need to include any children that right has as well
		node.setRight(right.getLeft());
		if(right.getLeft()!=null) right.getLeft().setParent(node);

		// demote node to right child's left position
		right.setLeft(node);

		// correct the parent reference
		right.setParent(node.getParent());

		// node points to new parent
		node.setParent(right);
	}

	/**
	 * The tree structure looks like a linked list from node with two children
	 * to the left and no right children Need to move node to right child of
	 * node's left child and make node's parent point to left
	 *
	 * @param node - the node about which to perform the rotation
	 */
	private void singleRightRotation(Node node) {
		Node left = node.getLeft();
		if (node.getParent() != null && node.getKey() < node.getParent().getKey()) {
			// node is the left child of the parent
			node.getParent().setLeft(left);
		} else if (node.getParent() != null) {
			// node is the parent's right child, set that reference instead
			node.getParent().setRight(left);
		}
		//need to include any children of left and move them to node
		node.setLeft(left.getRight());
		if(left.getRight()!=null) left.getRight().setParent(node);

		// demote node to left child's right position
		left.setRight(node);

		// correct the parent reference
		left.setParent(node.getParent());

		// node points to new parent
		node.setParent(left);
	}

	/**
	 * Need to rotate children up and left and parent down and right
	 *
	 * @param node - the node about which to rotate items
	 */
	private void leftRightRotation(Node node) {

		Node left = node.getLeft();
		Node newHead = left.getRight(); // this can't be null if we have
										// correctly determined we need to do a
										// LR rotation
		Node parent = node.getParent();

		newHead.setParent(parent);
		if (parent != null)
		{
			if (newHead.getKey() < parent.getKey())
			{
				parent.setLeft(newHead);
			}
			else
			{
				parent.setRight(newHead);
			}
		}

		Node t2 = newHead.getLeft();
		left.setRight(t2);
		if (t2 != null)
		{
			t2.setParent(left);
		}

		Node t3 = newHead.getRight();
		node.setLeft(t3);
		if (t3 != null)
		{
			t3.setParent(node);
		}

		newHead.setLeft(left);
		left.setParent(newHead);

		newHead.setRight(node);
		node.setParent(newHead);
	}

	/**
	 * Need to rotate children up and right and parent down and left
	 *
	 * @param node
	 *            - the node about which to rotate
	 */
	private void rightLeftRotation(Node node) {
		Node right = node.getRight();
		Node newHead = right.getLeft();
		Node parent = node.getParent();

		newHead.setParent(parent);
		if (parent != null)
		{
			if (newHead.getKey() < parent.getKey())
			{
				parent.setLeft(newHead);
			}
			else
			{
				parent.setRight(newHead);
			}
		}

		Node t2 = newHead.getLeft();
		node.setRight(t2);
		if (t2 != null)
		{
			t2.setParent(node);
		}

		Node t3 = newHead.getRight();
		right.setLeft(t3);
		if (t3 != null)
		{
			t3.setParent(right);
		}

		newHead.setLeft(node);
		node.setParent(newHead);

		newHead.setRight(right);
		right.setParent(newHead);
	}

	/**
	 * Tests whether or not the tree conforms to the structure of Left < Parent < Right
	 * and throws an exception if the state is violated
	 * @param node
	 */
	protected void isValidTree(Node node)
	{
		if(!Tester.DEBUG_MODE) return;
		
		int key = node.getKey();
		boolean valid = true;
		boolean balanced = true;
		if (node.getLeft() != null)
		{
			valid = node.getLeft().getKey() < key;
			balanced = Math.abs(node.getBalanceFactor()) < 2;
			isValidTree(node.getLeft());
		}
		if (node.getRight() != null)
		{
			valid = node.getRight().getKey() > key;
			balanced = (Math.abs(node.getBalanceFactor()) < 2);
			isValidTree(node.getRight());
		}
		if (!valid)
		{
			Node test = node.getParent() != null ? node.getParent() : node;
			if (test.getHeight() < 10)
			{
				printTreeLevel(Collections.singletonList(test), 0, test.getHeight());
			}
			throw new IllegalStateException("The tree is not in a valid state at node [" + node.getKey() + "]\n" +
					"items added in this order " + itemsAdded.toString());
		}
		else if (!balanced)
		{
			extraBalance++;
			rebalance(node);
			isValidTree(root);
		}
	}

	/**
	 * Used to print the tree to the console.
	 */
	public void printTree() {
		if(getRoot() == null) {
			System.err.println("Tree is empty");
			return;
		}

		int maxLevel = getRoot().getHeight();

		printTreeLevel(Collections.singletonList(getRoot()), 0, maxLevel);

	}

	/**
	 * A recursive function used to print to the console.
	 *
	 * @param tmpLevel
	 *            - the node representing the temporary root passed as a list -
	 *            passed as an array so to retain all left and right nodes of
	 *            the current level
	 * @param curLevel
	 *            - the current level that the algorithm is operating with
	 * @param maxLevel
	 *            - the maximum level of the tree
	 */
	private void printTreeLevel(List<Node> tmpLevel, int curLevel, int maxLevel) {

		// determine whether the array is empty
		if (tmpLevel.isEmpty())
			return;
		// determine whether all nodes in the array is empty
		// null nodes are entered to keep with formatting/spacing
		boolean allNull = true;
		for (Node node : tmpLevel) {
			if (node != null) {
				allNull = false;
			}
		}
		if (allNull)
			return;

		int floor = maxLevel - curLevel;
		int edgeLines = (int) Math.pow(2, (Math.max(floor - 1, 0)));
		int leftSpaces = (int) Math.pow(2, (floor)) - 1;
		int betweenSpaces = (int) Math.pow(2, (floor + 1)) - 1;

		printSpaces(leftSpaces); // print the left-most set of spaces

		// print all values within the array and add() sub-nodes
		List<Node> newNodes = new ArrayList<Node>();
		for (Node node : tmpLevel) {
			if (node != null) {
				System.err.print(node.getKey());
				newNodes.add(node.getLeft());
				newNodes.add(node.getRight());
			} else {
				newNodes.add(null);
				newNodes.add(null);
				System.err.print(" ");
			}
			printSpaces(betweenSpaces);
		}
		System.err.println("");

		// print the backslash and forwardslash for the tree's branches
		for (int i = 1; i <= edgeLines; i++) {
			for (int j = 0; j < tmpLevel.size(); j++) {
				printSpaces(leftSpaces - i);
				if (tmpLevel.get(j) == null) {
					printSpaces(edgeLines + edgeLines + i + 1);
					continue;
				}

				if (tmpLevel.get(j).getLeft() != null) {
					System.err.print("/");
				} else
					printSpaces(1);

				printSpaces(i + i - 1);

				if (tmpLevel.get(j).getRight() != null) {
					System.err.print("\\");
				} else
					printSpaces(1);

				printSpaces(edgeLines + edgeLines - i);
			}

			System.err.println("");
		}
		printTreeLevel(newNodes, curLevel + 1, maxLevel);
	}

	/**
	 * A helper function used to print white spaces to the console.
	 */
	private void printSpaces(int spaces) {
		for (int i = 0; i < spaces; i++)
			System.err.print(" ");
	}
}
